from pprint import pprint
from nicegui import ui, app
import asyncio
import rclpy
from rclpy.node import Node
from rclpy.executors import MultiThreadedExecutor
from unitree_go.msg import LowState, SportModeState
import threading
import math
import os
from go2_gait_planner.msg import GaitParam
from std_msgs.msg import Int32

from fastapi import Request

# 全局状态结构更新
robot_status = {
    "position": {"x": 0, "y": 0, "z": 0},
    "sim_position": {"x": 0, "y": 0, "z": 0},
    "target_position": {"x": 0, "y": 0, "z": 0},
    "joints": [0.0] * 12,
    "gait_params": {
        'gait_type': 0, 'velocity': 0.0, 'nominal_height': 0.25,
        'stance_duration': 0.2, 'stance_depth': 0.0, 'swing_height': 0.07
    },
    "lowstate": {
        "head": [0xFE, 0xEF],
        "level_flag": 0,
        "frame_reserve": 0,
        "sn": [0, 0],
        "version": [0, 0],
        "bandwidth": 0,
        "imu_state": {
            "quaternion": [0, 0, 0, 1],
            "gyroscope": [0, 0, 0],
            "accelerometer": [0, 0, 0],
            "rpy": [0, 0, 0],
            "temperature": 0
        },
        "motor_state": [{
            "mode": 0,
            "q": 0.0,
            "dq": 0.0,
            "ddq": 0.0,
            "tau_est": 0.0,
            "q_raw": 0.0,
            "dq_raw": 0.0,
            "ddq_raw": 0.0,
            "temperature": 0,
            "lost": 0,
            "reserve": [0, 0]
        } for _ in range(20)],
        "bms_state": {
            "version_high": 0,
            "version_low": 0,
            "status": 0,
            "soc": 0,
            "current": 0,
            "cycle": 0,
            "bq_ntc": [0, 0],
            "mcu_ntc": [0, 0],
            "cell_vol": [0] * 15
        },
        "foot_force": [0] * 4,
        "foot_force_est": [0] * 4,
        "tick": 0,
        "wireless_remote": [0] * 40,
        "bit_flag": 0,
        "adc_reel": 0.0,
        "temperature_ntc1": 0,
        "temperature_ntc2": 0,
        "power_v": 0.0,
        "power_a": 0.0,
        "fan_frequency": [0] * 4,
        "reserve": 0,
        "crc": 0
    },
    "movement_state": 0,
    "sport_mode_exists": False,
    "joystick_data": {"x": 0, "y": 0, "force": 0, "angle": 0},
    "current_rotation": 0  # 新增当前朝向角度
}

# ROS2配置
ROS_NODE_NAME = "nicegui_go2_controller"
GAIT_MSG_TOPIC = "/go2_gait_planner/gait_msg"
JOINTS_TOPIC = "/go2_gait_planner/joints"
LOWSTATE_TOPIC = "/lowstate"
SPORT_MODE_TOPIC = "/sportmodestate"
COMMAND_SERVICE = "/go2_command"
DEFAULT_VIDEO_PATH = r"/home/ljh/桌面/New/2.mp4"
VIDEO_WS_URL = "ws://localhost:8000/ws/video"


class GaitSubscriber(Node):
    def __init__(self):
        super().__init__('gait_subscriber')

        # 声明可调参数
        self.declare_parameter('gait_type', 0)
        self.declare_parameter('velocity', 0.5)
        self.declare_parameter('nominal_height', 0.25)
        self.declare_parameter('stance_duration', 0.2)
        self.declare_parameter('stance_depth', 0.0)
        self.declare_parameter('swing_height', 0.07)

        # 订阅器
        self.lowstate_sub = self.create_subscription(
            LowState, LOWSTATE_TOPIC, self.lowstate_callback, 10)
        self.subscription = self.create_subscription(
            GaitParam, GAIT_MSG_TOPIC, self.gait_callback, 10)
        self.joints_sub = self.create_subscription(
            GaitParam, JOINTS_TOPIC, self.joints_callback, 10)
        self.sport_mode_sub = self.create_subscription(
            SportModeState, SPORT_MODE_TOPIC, self.sport_mode_callback, 10)

    def gait_callback(self, msg):
        gait_params = {
            'gait_type': msg.gait_type,
            'velocity': msg.velocity,
            'nominal_height': msg.nominal_height,
            'stance_duration': msg.stance_duration,
            'stance_depth': msg.stance_depth,
            'swing_height': msg.swing_height,
            'movement': msg.movement
        }
        robot_status['gait_params'] = gait_params
        robot_status['movement_state'] = msg.movement
        self.get_logger().info(f"收到步态参数: {gait_params}")

    def joints_callback(self, msg):
        robot_status['joints'] = list(msg.data)

    def lowstate_callback(self, msg):
        """处理LowState消息"""
        robot_status['lowstate'] = {
            "head": list(msg.head),
            "level_flag": msg.level_flag,
            "frame_reserve": msg.frame_reserve,
            "sn": list(msg.sn),
            "version": list(msg.version),
            "bandwidth": msg.bandwidth,
            "imu_state": {
                "quaternion": list(msg.imu_state.quaternion),
                "gyroscope": list(msg.imu_state.gyroscope),
                "accelerometer": list(msg.imu_state.accelerometer),
                "rpy": list(msg.imu_state.rpy),
                "temperature": msg.imu_state.temperature
            },
            "motor_state": [{
                "mode": motor.mode,
                "q": motor.q,
                "dq": motor.dq,
                "ddq": motor.ddq,
                "tau_est": motor.tau_est,
                "q_raw": motor.q_raw,
                "dq_raw": motor.dq_raw,
                "ddq_raw": motor.ddq_raw,
                "temperature": motor.temperature,
                "lost": motor.lost,
                "reserve": list(motor.reserve)
            } for motor in msg.motor_state[:20]],
            "bms_state": {
                "version_high": msg.bms_state.version_high,
                "version_low": msg.bms_state.version_low,
                "status": msg.bms_state.status,
                "soc": msg.bms_state.soc,
                "current": msg.bms_state.current,
                "cycle": msg.bms_state.cycle,
                "bq_ntc": list(msg.bms_state.bq_ntc),
                "mcu_ntc": list(msg.bms_state.mcu_ntc),
                "cell_vol": list(msg.bms_state.cell_vol)
            },
            "foot_force": list(msg.foot_force),
            "foot_force_est": list(msg.foot_force_est),
            "tick": msg.tick,
            "wireless_remote": list(msg.wireless_remote),
            "bit_flag": msg.bit_flag,
            "adc_reel": msg.adc_reel,
            "temperature_ntc1": msg.temperature_ntc1,
            "temperature_ntc2": msg.temperature_ntc2,
            "power_v": msg.power_v,
            "power_a": msg.power_a,
            "fan_frequency": list(msg.fan_frequency),
            "reserve": msg.reserve,
            "crc": msg.crc
        }
        self.get_logger().info(f"收到LowState数据: tick={msg.tick}")

    def sport_mode_callback(self, msg):
        """处理SportModeState消息"""
        robot_status['position'] = {
            "x": msg.position[0],
            "y": msg.position[1],
            "z": msg.position[2]
        }
        robot_status['sport_mode_exists'] = True
        self.get_logger().info(
            f"收到SportModeState数据: position=({msg.position[0]:.2f}, {msg.position[1]:.2f}, {msg.position[2]:.2f})")


def ros_thread():
    try:
        rclpy.init()
        executor = MultiThreadedExecutor()
        gait_subscriber = GaitSubscriber()
        executor.add_node(gait_subscriber)
        executor.spin()
    except Exception as e:
        print(f"ROS线程错误: {str(e)}")
    finally:
        rclpy.shutdown()


@ui.page('/')
async def index():
    global status_label, video_container, lowstate_table, joystick_status

    if 'video_source' not in app.storage.user:
        if os.path.exists(DEFAULT_VIDEO_PATH):
            app.storage.user['video_source'] = DEFAULT_VIDEO_PATH
            app.storage.user['auto_play'] = True
        else:
            app.storage.user['video_source'] = ""
            ui.notify(f"默认视频文件 {DEFAULT_VIDEO_PATH} 不存在", type='warning')

    # 主界面布局 - 使用3行网格布局
    with ui.grid(rows=3).classes("w-full h-screen gap-4 p-4"):
        # 第一行：视频区域 (2列)
        with ui.grid(columns=2).classes("w-full h-[30vh] gap-4"):
            # 左边：推拉流视频
            with ui.card().classes("w-full h-full overflow-hidden"):
                ui.label("实时视频流").classes("text-xl")
                ui.add_body_html('''
                    <script src="https://cdn.jsdelivr.net/npm/flv.js@latest"></script>
                    <style>
                        #flv-container { width:100%; height:100% }
                        #flv-player { width:100%; height:100%; object-fit: contain }
                    </style>
                ''')
                video_container = ui.html('''
                    <div id="flv-container">
                        <video id="flv-player" controls muted></video>
                    </div>
                ''').classes("w-full h-[calc(100%-40px)]")
                ui.run_javascript('''
                    function initPlayer() {
                        const videoElement = document.getElementById('flv-player');
                        if (flvjs.isSupported()) {
                            const flvPlayer = flvjs.createPlayer({
                                type: 'flv',
                                url: 'http://127.0.0.1/live/test.live.flv',
                                isLive: true,
                                stashInitialSize: 128,
                                enableWorker: true,
                                lazyLoadMaxDuration: 3 * 60,
                                seekType: 'range'
                            });
                            flvPlayer.attachMediaElement(videoElement);
                            flvPlayer.load();
                            const playPromise = videoElement.play();
                            if (playPromise !== undefined) {
                                playPromise.catch(e => {
                                    console.log('自动播放被阻止，需要用户交互');
                                });
                            }
                            flvPlayer.on(flvjs.Events.ERROR, (errType, errDetail) => {
                                console.error('播放错误:', errType, errDetail);
                                if(errType === flvjs.ErrorTypes.NETWORK_ERROR) {
                                    setTimeout(() => {
                                        flvPlayer.unload();
                                        flvPlayer.detachMediaElement();
                                        flvPlayer.attachMediaElement(videoElement);
                                        flvPlayer.load();
                                        videoElement.play();
                                    }, 3000);
                                }
                            });
                            window.flvPlayer = flvPlayer;
                        } else {
                            console.error('浏览器不支持FLV播放');
                            videoElement.innerHTML = '您的浏览器不支持视频播放';
                        }
                    }
                    setTimeout(initPlayer, 500);
                ''')

            # 右边：本地视频
            with ui.card().classes("w-full h-full overflow-hidden"):
                ui.label("本地视频").classes("text-xl")
                video_player = ui.video("").classes("w-full h-[calc(100%-40px)]").props("controls autoplay muted")
                if app.storage.user['video_source']:
                    video_player.set_source(DEFAULT_VIDEO_PATH)

        # 第二行：控制区域 (2列)
        with ui.grid(columns=2).classes("w-full h-[20vh] gap-4"):
            # 左边：控制按钮
            with ui.card().classes("w-full h-full overflow-auto"):
                ui.label("机器人控制").classes("text-xl")
                status_label = ui.label(f"状态: 位置(0.00, 0.00)").classes("text-lg")
                with ui.row().classes("w-full gap-2"):
                    ui.button("前进", on_click=lambda: set_movement_state(1)).classes("bg-green-500 hover:bg-green-600")
                    ui.button("后退", on_click=lambda: set_movement_state(2)).classes("bg-red-500 hover:bg-red-600")
                    ui.button("左移", on_click=lambda: set_movement_state(3)).classes("bg-blue-500 hover:bg-blue-600")
                    ui.button("右移", on_click=lambda: set_movement_state(4)).classes("bg-blue-500 hover:bg-blue-600")
                    ui.button("停止", on_click=lambda: set_movement_state(0)).classes("bg-gray-500 hover:bg-gray-600")

            # 右边：参数表
            with ui.card().classes("w-full h-full overflow-auto"):
                ui.label("LowState 参数表").classes("text-xl")
                lowstate_table = ui.table(
                    columns=[
                        {'name': 'group', 'label': '参数组', 'field': 'group', 'align': 'left'},
                        {'name': 'param', 'label': '参数', 'field': 'param', 'align': 'left'},
                        {'name': 'value', 'label': '值', 'field': 'value', 'align': 'right'}
                    ],
                    rows=[],
                    row_key='id'
                ).classes("w-full")

        # 第三行：3D场景 (全宽)
        with ui.card().classes("w-full h-[45vh] overflow-hidden relative"):  # 添加relative定位
            ui.label("机器人3D空间").classes("text-xl")

            # 摇杆状态显示 - 移出摇杆容器避免遮挡
            joystick_status = ui.label("摇杆: 待机").classes(
                'text-xs absolute top-2 right-2 text-gray-600 bg-white/70 px-2 py-1 rounded-full z-20')

            # 摇杆控制容器 - 固定在3D场景右下角
            with ui.element('div').classes('absolute bottom-4 right-4 z-10'):
                with ui.card().classes('w-40 h-40 bg-gray-100/80 rounded-full p-0 overflow-hidden'):
                    # 摇杆容器
                    joystick_container = ui.html(
                        '<div id="joystick-container" style="width:150px;height:150px;margin:0 auto;"></div>'
                    )

                    # 站立/趴下控制按钮
                    with ui.row().classes('absolute bottom-2 w-full justify-center gap-2'):
                        ui.button("站立", on_click=lambda: send_stand_command(1)) \
                            .classes('bg-green-500 hover:bg-green-600 text-xs py-1 px-2')
                        ui.button("趴下", on_click=lambda: send_stand_command(0)) \
                            .classes('bg-red-500 hover:bg-red-600 text-xs py-1 px-2')

            # 3D场景HTML
            ui.add_body_html('''
                <script src="https://cdn.jsdelivr.net/npm/three@0.132.2/build/three.min.js"></script>
                <script src="https://cdn.jsdelivr.net/npm/three@0.132.2/examples/js/controls/OrbitControls.min.js"></script>
                <script src="https://cdn.jsdelivr.net/npm/gsap@3.11.4/dist/gsap.min.js"></script>
                <script src="https://cdn.jsdelivr.net/npm/nipplejs@0.10.1/dist/nipplejs.min.js"></script>
            ''')
            scene_container = ui.html('''
                <div id="threejs-container" style="width:100%; height:100%; position:relative;">
                    <canvas id="threejs-canvas" style="display:block; width:100%; height:100%"></canvas>
                </div>
            ''').classes("w-full h-full")

            # 修正的3D场景初始化脚本
            ui.run_javascript('''
                // 初始化Three.js场景
                const container = document.getElementById('threejs-container');
                const canvas = document.getElementById('threejs-canvas');
                canvas.width = container.clientWidth;
                canvas.height = container.clientHeight;

                const scene = new THREE.Scene();
                scene.background = new THREE.Color(0xf0f0f0);

                // 修正相机位置和朝向
                const camera = new THREE.PerspectiveCamera(60, container.clientWidth / container.clientHeight, 0.1, 1000);
                camera.position.set(5, 5, 5);
                camera.lookAt(0, 1, 0);

                const renderer = new THREE.WebGLRenderer({ 
                    canvas: canvas, 
                    antialias: true,
                    alpha: true
                });
                renderer.setSize(canvas.width, canvas.height);
                renderer.setClearColor(0x000000, 0);

                // 添加辅助工具
                const axesHelper = new THREE.AxesHelper(3);
                scene.add(axesHelper);

                // 扩大网格尺寸和密度
                const gridHelper = new THREE.GridHelper(50, 50);  // 尺寸50m，50个单元格
                gridHelper.position.y = 0;
                scene.add(gridHelper);

                // 灯光 - 增加光源强度
                const ambientLight = new THREE.AmbientLight(0xffffff, 0.6);
                scene.add(ambientLight);

                const directionalLight = new THREE.DirectionalLight(0xffffff, 1);
                directionalLight.position.set(2, 3, 1);
                directionalLight.castShadow = true;
                scene.add(directionalLight);

                // 创建小狗模型
                const dog = new THREE.Group();
                scene.add(dog);

                // 身体 - 圆柱体
                const bodyGeometry = new THREE.CylinderGeometry(0.5, 0.5, 1.5, 16);
                bodyGeometry.rotateX(Math.PI/2);
                const body = new THREE.Mesh(
                    bodyGeometry,
                    new THREE.MeshPhongMaterial({ color: 0x8B4513 })
                );
                dog.add(body);

                // 头部 - 球体
                const head = new THREE.Mesh(
                    new THREE.SphereGeometry(0.4, 16, 16),
                    new THREE.MeshPhongMaterial({ color: 0x8B4513 })
                );
                head.position.set(0.8, 0.2, 0);
                dog.add(head);

                // 鼻子
                const nose = new THREE.Mesh(
                    new THREE.SphereGeometry(0.1, 16, 16),
                    new THREE.MeshPhongMaterial({ color: 0x000000 })
                );
                nose.position.set(1.1, 0.2, 0);
                dog.add(nose);

                // 耳朵
                const earGeometry = new THREE.ConeGeometry(0.2, 0.4, 16);
                const leftEar = new THREE.Mesh(
                    earGeometry,
                    new THREE.MeshPhongMaterial({ color: 0x8B4513 })
                );
                leftEar.position.set(0.6, 0.5, 0.3);
                leftEar.rotation.z = -Math.PI/4;
                dog.add(leftEar);

                const rightEar = new THREE.Mesh(
                    earGeometry,
                    new THREE.MeshPhongMaterial({ color: 0x8B4513 })
                );
                rightEar.position.set(0.6, 0.5, -0.3);
                rightEar.rotation.z = -Math.PI/4;
                dog.add(rightEar);

                // 眼睛
                const eyeGeometry = new THREE.SphereGeometry(0.05, 16, 16);
                const leftEye = new THREE.Mesh(
                    eyeGeometry,
                    new THREE.MeshPhongMaterial({ color: 0x000000 })
                );
                leftEye.position.set(0.9, 0.3, 0.15);
                dog.add(leftEye);

                const rightEye = new THREE.Mesh(
                    eyeGeometry,
                    new THREE.MeshPhongMaterial({ color: 0x000000 })
                );
                rightEye.position.set(0.9, 0.3, -0.15);
                dog.add(rightEye);

                // 腿
                const legGeometry = new THREE.CylinderGeometry(0.1, 0.1, 0.6, 8);
                const legPositions = [
                        { x: 0.3, y: -0.75, z: 0.3 },
                        { x: -0.3, y: -0.75, z: 0.3 },
                        { x: 0.3, y: -0.75, z: -0.3 },
                        { x: -0.3, y: -0.75, z: -0.3 }
                    ];

                const legs = [];
                legPositions.forEach(pos => {
                    const leg = new THREE.Mesh(
                        legGeometry,
                        new THREE.MeshPhongMaterial({ color: 0x8B4513 })
                    );
                    leg.position.set(pos.x, pos.y, pos.z);
                    dog.add(leg);
                    legs.push(leg);
                });

                // 尾巴
                const tailGeometry = new THREE.CylinderGeometry(0.05, 0.05, 0.4, 8);
                const tail = new THREE.Mesh(
                    tailGeometry,
                    new THREE.MeshPhongMaterial({ color: 0x8B4513 })
                );
                tail.position.set(-0.8, -0.2, 0);
                tail.rotation.z = Math.PI/4;
                dog.add(tail);

                // 方向指示器（箭头）
                const arrowGeometry = new THREE.ConeGeometry(0.2, 0.5, 8);
                const arrow = new THREE.Mesh(
                    arrowGeometry,
                    new THREE.MeshPhongMaterial({ color: 0xff0000 })
                );
                arrow.position.set(0, 0.5, 0);
                arrow.rotation.x = Math.PI/2;
                dog.add(arrow);

                // 控制器和动画循环
                const controls = new THREE.OrbitControls(camera, renderer.domElement);
                controls.enableDamping = true;
                controls.dampingFactor = 0.05;
                controls.screenSpacePanning = true;

                function animate() {
                    requestAnimationFrame(animate);
                    controls.update();
                    renderer.render(scene, camera);
                }
                animate();

                // 存储到全局
                window.robotScene = { scene, dog, legs, camera, renderer, controls, arrow };

                // 响应式调整
                window.addEventListener('resize', () => {
                    camera.aspect = container.clientWidth / container.clientHeight;
                    camera.updateProjectionMatrix();
                    renderer.setSize(container.clientWidth, container.clientHeight);
                });

                // 初始化摇杆 - 使用正确的容器ID
                setTimeout(() => {
                    const joystickZone = document.getElementById('joystick-container');
                    if (!joystickZone) {
                        console.error('摇杆容器未找到');
                        return;
                    }

                    const manager = nipplejs.create({
                        zone: joystickZone,
                        mode: 'static',
                        position: { left: '50%', top: '50%' },
                        color: '#3b82f6',
                        size: 120,
                        dynamicPage: true,
                        restJoystick: true,
                        restOpacity: 0.1,
                        lockY: false,
                        lockX: false
                    });

                    manager.on('start', function() {
                        console.log('摇杆开始控制');
                        pywebview.api.updateJoystickStatus('控制中');
                    });

                    manager.on('move', function(evt, data) {
                        const force = data.force || 0;
                        const angle = data.direction ? data.direction.angle : 'none';
                        const x = Math.round(data.vector.x * 100) / 100;
                        const y = Math.round(data.vector.y * 100) / 100;

                        // 发送摇杆数据到Python
                        fetch('/joystick', {
                            method: 'POST',
                            headers: {
                                'Content-Type': 'application/json'
                            },
                            body: JSON.stringify({
                                x: x,
                                y: y,
                                force: force,
                                angle: angle
                            })
                        }).catch(e => console.error('发送摇杆数据失败:', e));
                    });

                    manager.on('end', function() {
                        console.log('摇杆结束控制');
                        pywebview.api.updateJoystickStatus('待机');
                        fetch('/joystick', {
                            method: 'POST',
                            headers: {
                                'Content-Type': 'application/json'
                            },
                            body: JSON.stringify({ x: 0, y: 0, force: 0, angle: 'none' })
                        });
                    });

                    window.joystickManager = manager;
                }, 500);
            ''')

    # 定时更新
    ui.timer(0.05, update_robot_position)
    ui.timer(0.5, update_lowstate_table)
    await connect_video_stream()

    # 注册摇杆处理API
    @ui.post('/joystick')
    async def handle_joystick(request: Request):
        """处理摇杆输入"""
        data = await request.json()
        x = data.get('x', 0)
        y = data.get('y', 0)
        force = data.get('force', 0)
        angle = data.get('angle', 'none')

        # 更新全局摇杆数据
        robot_status['joystick_data'] = {
            "x": x,
            "y": y,
            "force": force,
            "angle": angle if angle != 'none' else 0
        }

        # 力度小于0.1视为无效输入
        if force < 0.1:
            send_gait_command(movement=0)
            return

        # 计算方向（角度以度为单位）
        if angle == 'none':
            return

        # 保存角度用于方向控制
        robot_status['joystick_data']['angle'] = float(angle)

        # 方向映射（连续控制）
        movement = 0
        velocity = max(0.3, min(1.0, force * 2.0))  # 增加速度系数

        # 发送运动指令（使用连续角度）
        send_gait_command(movement=movement, velocity=velocity)

    @ui.page('/joystick-status')
    async def update_joystick_status(status: str):
        """更新摇杆状态显示"""
        joystick_status.text = f"摇杆: {status}"


def set_movement_state(state):
    """设置机器人运动状态"""
    robot_status['movement_state'] = state
    if state == 0:
        send_gait_command(movement=0)
    else:
        send_gait_command(movement=state, velocity=0.8)  # 增加默认速度


def update_robot_position():
    """更新机器人位置，实现平滑移动"""
    # 如果SportModeState节点存在，使用真实位置
    if robot_status['sport_mode_exists']:
        pos = robot_status['position']
        movement_state = robot_status['movement_state']
        status_label.text = f"状态: 真实位置({pos['x']:.2f}, {pos['y']:.2f}) 运动:{['停止', '前进', '后退', '左移', '右移'][movement_state]}"

        # 更新3D场景使用真实位置
        ui.run_javascript(f'''
            if (window.robotScene) {{
                const dog = window.robotScene.dog;
                const legs = window.robotScene.legs;
                const arrow = window.robotScene.arrow;

                // 使用GSAP实现平滑动画
                gsap.to(dog.position, {{
                    x: {pos['x']},
                    z: {pos['y']},  
                    y: {pos['z']},  
                    duration: 0.2,
                    ease: "power2.out"
                }});

                // 根据运动状态调整小狗朝向
                if ({movement_state} > 0) {{
                    const targetRotation = {{
                        1: Math.PI,    // 前进
                        2: 0,          // 后退
                        3: Math.PI/2,  // 左移
                        4: -Math.PI/2  // 右移
                    }}[{movement_state}];

                    // 更新全局旋转状态
                    pywebview.api.updateRotation(targetRotation);

                    gsap.to(dog.rotation, {{
                        y: targetRotation,
                        duration: 0.3,
                        ease: "power2.out"
                    }});
                }}

                // 模拟腿部动画（增强幅度）
                const time = Date.now() * 0.001;
                if ({movement_state} > 0) {{
                    legs.forEach((leg, i) => {{
                        const legAngle = Math.sin(time * 6 + i * Math.PI/2) * 0.4;  // 增加摆动幅度
                        leg.rotation.z = legAngle;
                    }});
                }} else {{
                    legs.forEach(leg => {{
                        leg.rotation.z = 0;
                    }});
                }}
            }}
        ''')
    else:
        # SportModeState节点不存在，使用仿真位置
        pos = robot_status['sim_position']
        target = robot_status['target_position']
        movement_state = robot_status['movement_state']
        joystick = robot_status['joystick_data']

        # 获取当前旋转角度
        current_rotation = robot_status.get('current_rotation', 0)

        # 新增：如果使用摇杆控制，根据摇杆数据更新目标位置
        if joystick['force'] > 0.1:
            # 摇杆控制优先
            speed_factor = 0.15 * joystick['force']  # 增加移动速度

            # 根据当前朝向旋转向量
            angle_rad = current_rotation
            cos_angle = math.cos(angle_rad)
            sin_angle = math.sin(angle_rad)

            # 旋转摇杆向量到世界坐标系
            dx = joystick['x'] * cos_angle - joystick['y'] * sin_angle
            dy = joystick['x'] * sin_angle + joystick['y'] * cos_angle

            target['x'] += dx * speed_factor
            target['y'] += dy * speed_factor

            # 更新旋转角度（如果摇杆有角度）
            if joystick['angle'] != 0:
                robot_status['current_rotation'] = math.radians(joystick['angle'])
        else:
            # 原有按钮控制逻辑
            speed = 0.08  # 增加按钮控制速度
            if movement_state == 1:  # 前进
                target['y'] += speed * math.cos(current_rotation)
                target['x'] += speed * math.sin(current_rotation)
            elif movement_state == 2:  # 后退
                target['y'] -= speed * math.cos(current_rotation)
                target['x'] -= speed * math.sin(current_rotation)
            elif movement_state == 3:  # 左移
                target['x'] -= speed * math.cos(current_rotation)
                target['y'] += speed * math.sin(current_rotation)
            elif movement_state == 4:  # 右移
                target['x'] += speed * math.cos(current_rotation)
                target['y'] -= speed * math.sin(current_rotation)

        # 使用线性插值实现平滑移动
        pos['x'] += (target['x'] - pos['x']) * 0.1
        pos['y'] += (target['y'] - pos['y']) * 0.1

        status_label.text = f"状态: 仿真位置({pos['x']:.2f}, {pos['y']:.2f}) 运动:{['停止', '前进', '后退', '左移', '右移'][movement_state]}"

        # 更新3D场景使用仿真位置
        ui.run_javascript(f'''
            if (window.robotScene) {{
                const dog = window.robotScene.dog;
                const legs = window.robotScene.legs;
                const arrow = window.robotScene.arrow;

                // 使用GSAP实现平滑动画
                gsap.to(dog.position, {{
                    x: {pos['x']},
                    z: {pos['y']},
                    duration: 0.2,
                    ease: "power2.out"
                }});

                // 根据运动状态调整小狗朝向
                if ({movement_state} > 0) {{
                    let targetRotation = 0;
                    switch({movement_state}) {{
                        case 1: targetRotation = Math.PI; break;    // 前进
                        case 2: targetRotation = 0; break;          // 后退
                        case 3: targetRotation = Math.PI/2; break;  // 左移
                        case 4: targetRotation = -Math.PI/2; break; // 右移
                    }}

                    // 更新全局旋转状态
                    pywebview.api.updateRotation(targetRotation);

                    gsap.to(dog.rotation, {{
                        y: targetRotation,
                        duration: 0.3,
                        ease: "power2.out"
                    }});
                }}
                else if ({joystick['force']} > 0.1) {{
                    // 摇杆控制时更新朝向
                    const targetRotation = {joystick['angle']} * Math.PI / 180;
                    pywebview.api.updateRotation(targetRotation);

                    gsap.to(dog.rotation, {{
                        y: targetRotation,
                        duration: 0.3,
                        ease: "power2.out"
                    }});
                }}

                // 模拟腿部动画（增强幅度）
                const time = Date.now() * 0.001;
                if ({movement_state} > 0 || {joystick['force']} > 0.1) {{
                    legs.forEach((leg, i) => {{
                        const legAngle = Math.sin(time * 6 + i * Math.PI/2) * 0.4;  // 增加摆动幅度
                        leg.rotation.z = legAngle;
                    }});
                }} else {{
                    legs.forEach(leg => {{
                        leg.rotation.z = 0;
                    }});
                }}

                // 更新箭头方向
                arrow.rotation.y = dog.rotation.y;
            }}
        ''')


def update_lowstate_table():
    """更新/lowstate参数表"""
    lowstate = robot_status['lowstate']
    gait_params = robot_status['gait_params']

    rows = [
        # 系统信息
        {'id': 'sys1', 'group': '系统', 'param': '运行时长(ms)', 'value': str(lowstate['tick'])},
        {'id': 'sys2', 'group': '系统', 'param': '电池电量(%)', 'value': f"{lowstate['bms_state']['soc']}%"},
        {'id': 'sys3', 'group': '系统', 'param': '位置源',
         'value': "真实位置" if robot_status['sport_mode_exists'] else "仿真位置"},

        # IMU状态
        {'id': 'imu1', 'group': 'IMU', 'param': '四元数(W)', 'value': f"{lowstate['imu_state']['quaternion'][3]:.3f}"},
        {'id': 'imu2', 'group': 'IMU', 'param': 'RPY(deg)',
         'value': f"{math.degrees(lowstate['imu_state']['rpy'][0]):.1f}°, {math.degrees(lowstate['imu_state']['rpy'][1]):.1f}°, {math.degrees(lowstate['imu_state']['rpy'][2]):.1f}°"},

        # 电机状态（示例显示前腿两个电机）
        {'id': 'm1', 'group': '电机FR', 'param': '角度(rad)', 'value': f"{lowstate['motor_state'][0]['q']:.3f}"},
        {'id': 'm2', 'group': '电机FR', 'param': '扭矩(Nm)', 'value': f"{lowstate['motor_state'][0]['tau_est']:.2f}"},
        {'id': 'm3', 'group': '电机FL', 'param': '角度(rad)', 'value': f"{lowstate['motor_state'][1]['q']:.3f}"},

        # 电源状态
        {'id': 'pwr1', 'group': '电源', 'param': '电压(V)', 'value': f"{lowstate['power_v']:.1f}"},
        {'id': 'pwr2', 'group': '电源', 'param': '电流(A)', 'value': f"{lowstate['power_a']:.1f}"},

        # 步态参数
        {'id': 'g1', 'group': '步态', 'param': '运动状态',
         'value': ['停止', '前进', '后退', '左移', '右移'][gait_params.get('movement', 0)]},
        {'id': 'g2', 'group': '步态', 'param': '速度(m/s)', 'value': f"{gait_params['velocity']:.2f}"}
    ]

    # 添加足端力传感器数据
    for i in range(4):
        rows.append({
            'id': f'foot{i}',
            'group': f'足端力{i + 1}',
            'param': '测量值(N)',
            'value': f"{lowstate['foot_force'][i]}"
        })

    # 如果SportModeState节点存在，添加真实位置信息
    if robot_status['sport_mode_exists']:
        pos = robot_status['position']
        rows.append({
            'id': 'pos1',
            'group': '位置',
            'param': 'X坐标(m)',
            'value': f"{pos['x']:.2f}"
        })
        rows.append({
            'id': 'pos2',
            'group': '位置',
            'param': 'Y坐标(m)',
            'value': f"{pos['y']:.2f}"
        })
        rows.append({
            'id': 'pos3',
            'group': '位置',
            'param': 'Z坐标(m)',
            'value': f"{pos['z']:.2f}"
        })

    lowstate_table.rows = rows
    lowstate_table.update()


async def connect_video_stream():
    global video_ws
    try:
        video_ws = await ui.run_javascript(f"""
            const socket = new WebSocket('{VIDEO_WS_URL}');
            socket.onopen = () => setInterval(() => socket.send('ping'), 15000);
            return socket;
        """)
    except Exception as e:
        print(f"连接视频流失败: {e}")


def send_gait_command(velocity=1.0, nominal_height=0.25, stance_duration=0.2,
                      stance_depth=0.0, swing_height=0.07, gait_type=0, movement=0):
    try:
        node = rclpy.create_node('gait_command_publisher')
        pub = node.create_publisher(GaitParam, GAIT_MSG_TOPIC, 10)

        gait_msg = GaitParam()
        gait_msg.gait_type = gait_type
        gait_msg.velocity = velocity
        gait_msg.nominal_height = nominal_height
        gait_msg.stance_duration = stance_duration
        gait_msg.stance_depth = stance_depth
        gait_msg.swing_height = swing_height
        gait_msg.movement = movement

        pub.publish(gait_msg)
        node.destroy_node()

        robot_status['gait_params'] = {
            'gait_type': gait_type,
            'velocity': velocity,
            'nominal_height': nominal_height,
            'stance_duration': stance_duration,
            'stance_depth': stance_depth,
            'swing_height': swing_height,
            'movement': movement
        }
        ui.notify(f"已发送运动指令: {'停止' if movement == 0 else ['前进', '后退', '左移', '右移'][movement - 1]}")
    except Exception as e:
        ui.notify(f"ROS2 Error: {str(e)}", type='negative')


def send_stand_command(stand: int):
    """发送站立/趴下命令"""
    try:
        node = rclpy.create_node('stand_command_publisher')
        pub = node.create_publisher(Int32, "/go2_gait_planner/stand_sit", 10)

        msg = Int32()
        msg.data = stand
        pub.publish(msg)

        node.destroy_node()
        ui.notify(f"已发送{'站立' if stand == 1 else '趴下'}命令")
    except Exception as e:
        ui.notify(f"ROS2 Error: {str(e)}", type='negative')


# 启动ROS2线程
threading.Thread(target=ros_thread, daemon=True).start()

# 启动NiceGUI应用
ui.run(
    title="Unitree Go2 Control Panel",
    port=5000,
    storage_secret="111",
    dark=False,
    host='0.0.0.0'
)